package com.wpml.netty.nio.multi;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.channels.Channel;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description 线程组管理
 * @Date 2021-07-17
 * @Author wangpei
 **/
public class WorkerGroup {
    // 多路复用器
    Selector selector;
    // 工作线程
    WorkerThread[] workers;
    // 主从配置使用
    WorkerGroup workerGroup;

    ServerSocketChannel server;
    // 计数器
    AtomicInteger xid = new AtomicInteger(0);

    /**
     * 构造中初始化 selector 和 工作线程
     **/
    public WorkerGroup(int threadNum) {
        try {
            selector = Selector.open();
            workers = new WorkerThread[threadNum];

            for (int i = 0; i < threadNum; i++) {
                // 初始化线程
                workers[i] = new WorkerThread();
                // 启动线程
                new Thread(workers[i]).start();
            }
        } catch (IOException e) {
            e.printStackTrace();
            System.exit(1);
        }
    }


    public void setWorkerGroup(WorkerGroup worker) {
        this.workerGroup = worker;
    }

    /**
     * 绑定服务端口， 选择一个线程，注册selector
     */
    public void bind(int port) {
        try {
            server = ServerSocketChannel.open();
            // 设置非阻塞
            server.configureBlocking(false);
            server.bind(new InetSocketAddress(port));

            selectWorkThread(server);

        } catch (IOException | InterruptedException e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    /**
     * 选择工作线程注册监听
     * @param ch  server端或client端
     */
    public void selectWorkThread(Channel ch) throws InterruptedException {
        if (ch instanceof ServerSocketChannel) {
            WorkerThread serverThread = selectOne();
            serverThread.lbq.put(ch);
            serverThread.setWorkerGroup(workerGroup);
            // 唤醒阻塞着的selector
            serverThread.selector.wakeup();
        }
        if (ch instanceof SocketChannel) {
            WorkerThread clientThread = selectOne();
            clientThread.lbq.put(ch);
            clientThread.selector.wakeup();
        }
    }

    /**
     * 选择一个工作线程 workerThread
     * @return
     */
    private WorkerThread selectOne() {
        // 轮训取出一个worker
        int index = xid.getAndIncrement() % workers.length;
        return workers[index];
    }
}
